{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n",
      "13\n",
      "11\n",
      "3\n",
      "14\n",
      "14\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import  os \n",
    "from os.path import join, isdir, isfile, exists\n",
    "from os import makedirs\n",
    "\n",
    "import cv2\n",
    "from glob import glob\n",
    "\n",
    "import json\n",
    "\n",
    "import glob\n",
    "\n",
    "import numpy as np\n",
    "def makedir(path):\n",
    "    \n",
    "    if not exists(path):\n",
    "        makedirs(path)\n",
    "import cv2\n",
    "\n",
    "import cv2\n",
    "\n",
    "\n",
    "'''\n",
    "description:  convert data from PR_ant_bee dataset format to TOPIC_dataset_format \n",
    "return {*}\n",
    "'''\n",
    "class BeeDatasetModifier:\n",
    "\n",
    "    def __init__(self,path):\n",
    "\n",
    "        # subsets =['test','train',]\n",
    "        subsets =['test']\n",
    "        self.dance_bee_id = {\n",
    "            'bee02':13,\n",
    "            'bee03':11,\n",
    "            'bee05':3,\n",
    "            'bee06':14,\n",
    "            'bee07':14,\n",
    "        }\n",
    "\n",
    "        all_seqs = [] \n",
    "        for subset_name in subsets :\n",
    "            seqs = glob.glob(join(path,subset_name)+'/*')\n",
    "\n",
    "            all_seqs += seqs\n",
    "        self.all_seqs = sorted(all_seqs)\n",
    "\n",
    "    def __len__(self):\n",
    "\n",
    "        return len(self.all_seqs)\n",
    "\n",
    "    \n",
    "\n",
    "    '''\n",
    "    description:  only keep the target bee \n",
    "    param {*} self\n",
    "    param {*} target_id\n",
    "    return {*}\n",
    "    '''\n",
    "    def modify_gt(self,idx, target_id):\n",
    "\n",
    "        _,labels = converter.getitem(idx)\n",
    "\n",
    "        \n",
    "        new_labels = np.array([line  for line in labels if line[1] == target_id])\n",
    "\n",
    "        path = self.all_seqs[idx]\n",
    "        np.savetxt(join(path,'gt','gt.txt'),new_labels,delimiter=',',fmt=\"%d\")\n",
    "\n",
    "        \n",
    "        return new_labels\n",
    "\n",
    "    def modify_all(self,):\n",
    "        self.dance_bee_id\n",
    "\n",
    "        \n",
    "\n",
    "        for idx in range(self.__len__()):\n",
    "            key = self.all_seqs[idx].split('/')[-1].split('.')[0]\n",
    "            \n",
    "            print(self.dance_bee_id[key])\n",
    "            self.modify_gt(idx,self.dance_bee_id[key])\n",
    "\n",
    "            \n",
    "\n",
    "        \n",
    "\n",
    "    def getitem(self,idx):\n",
    "\n",
    "        path = self.all_seqs[idx]\n",
    "\n",
    "\n",
    "        labels = np.loadtxt(join(path,'gt','gt.txt'),delimiter=',')\n",
    "\n",
    "        img_names = sorted(glob.glob(join(path,'img1','*.jpg')))\n",
    "\n",
    "        return img_names,labels\n",
    "\n",
    "\n",
    "        \n",
    "root = \"/data/xusc/exp/topictrack-bee/data/beedance\"\n",
    "converter = BeeDatasetModifier(root)\n",
    "\n",
    "\n",
    "#* target id = 11 \n",
    "\n",
    "\n",
    "print(converter.__len__())\n",
    "\n",
    "converter.modify_all()\n",
    "# new_labels = converter.modify_gt(0,11)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/data/xusc/exp/topictrack-bee/data/beedance/test/bee02',\n",
       " '/data/xusc/exp/topictrack-bee/data/beedance/test/bee03',\n",
       " '/data/xusc/exp/topictrack-bee/data/beedance/test/bee05',\n",
       " '/data/xusc/exp/topictrack-bee/data/beedance/test/bee06',\n",
       " '/data/xusc/exp/topictrack-bee/data/beedance/test/bee07']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "converter.all_seqs\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 19.,  11., 468.,   5., 397., 109.,   1.,   1.,   1.],\n",
       "       [ 20.,  11., 478.,   4., 397., 109.,   1.,   1.,   1.],\n",
       "       [ 21.,  11., 468.,  26., 397., 109.,   1.,   1.,   1.],\n",
       "       [ 22.,  11., 474.,  30., 397., 109.,   1.,   1.,   1.],\n",
       "       [ 23.,  11., 445.,  -5., 429., 148.,   1.,   1.,   1.],\n",
       "       [ 24.,  11., 435.,   0., 429., 148.,   1.,   1.,   1.],\n",
       "       [ 25.,  11., 432.,  -4., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 26.,  11., 425.,   2., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 27.,  11., 423.,   1., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 28.,  11., 421.,   1., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 29.,  11., 419.,   5., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 30.,  11., 419.,   8., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 31.,  11., 421.,   2., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 32.,  11., 411.,   8., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 33.,  11., 426.,   7., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 34.,  11., 424.,  10., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 35.,  11., 423.,   7., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 36.,  11., 424.,   4., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 37.,  11., 422.,   1., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 38.,  11., 415.,   0., 450., 155.,   1.,   1.,   1.],\n",
       "       [ 39.,  11., 402.,  -1., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 40.,  11., 402.,   1., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 41.,  11., 390.,   1., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 42.,  11., 395.,   0., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 43.,  11., 396.,   0., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 44.,  11., 444., -34., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 45.,  11., 435., -33., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 46.,  11., 420.,   4., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 47.,  11., 436.,  22., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 48.,  11., 441.,  23., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 49.,  11., 432.,  19., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 50.,  11., 420.,   3., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 51.,  11., 395.,   8., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 52.,  11., 393.,  14., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 53.,  11., 403.,   9., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 54.,  11., 410.,   5., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 55.,  11., 414.,   1., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 56.,  11., 443.,  16., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 57.,  11., 455.,  20., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 58.,  11., 465.,  20., 472., 163.,   1.,   1.,   1.],\n",
       "       [ 59.,  11., 470.,  24., 472., 163.,   1.,   1.,   1.]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "converter.modify_gt(0,11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_labels = converter.modify_gt(0,11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  1.,  11., 257., ...,   1.,   1.,   1.],\n",
       "       [  2.,  11., 271., ...,   1.,   1.,   1.],\n",
       "       [  3.,  11., 166., ...,   1.,   1.,   1.],\n",
       "       ...,\n",
       "       [172.,  11., 416., ...,   1.,   1.,   1.],\n",
       "       [173.,  11., 429., ...,   1.,   1.,   1.],\n",
       "       [174.,  11., 588., ...,   1.,   1.,   1.]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yang_real",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
